from flask import Flask, request, jsonify
from typing import List, Dict, Optional, Tuple
import os
import json
import requests
import logging
import subprocess
from config import Config
from db_handler import DatabaseHandler

logger = logging.getLogger(__name__)

app = Flask(__name__)

class PDFImageProcessor:
    @staticmethod
    def create_with_api_key(api_key: str):
        return PDFImageProcessor(api_key=api_key)
        
    def __init__(self, api_key: str):
        """
        初始化 PDF 图片处理器
        
        Args:
            api_key: API 认证密钥
        """
        self.api_key = api_key
        self.headers = {"Authorization": f"Bearer {api_key}"}

    def extract_images_from_pdfs(self, input_dir: str, output_dir: str) -> None:
        """
        从PDF文件中提取图片
        
        Args:
            input_dir: 输入目录
            output_dir: 输出目录
        """
        # 创建临时目录
        os.makedirs(output_dir, exist_ok=True)
        
        # 遍历输入目录中的PDF文件
        for filename in os.listdir(input_dir):
            if filename.lower().endswith('.pdf'):
                pdf_path = os.path.join(input_dir, filename)
                pdf_name = os.path.splitext(filename)[0]
                out_path = os.path.join(output_dir, pdf_name)
                os.makedirs(out_path, exist_ok=True)

                # 使用 magic-pdf 提取图片
                cmd = ['magic-pdf', '-p', pdf_path, '-o', out_path]
                logger.info(f"正在解析 PDF: {filename}")
                try:
                    subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
                    logger.info(f"解析成功: {filename}")
                except subprocess.CalledProcessError as e:
                    logger.error(f"解析失败: {filename}")
                    logger.error(e.stderr)
                    continue

                # 上传图片到 MinIO
                images_path = os.path.join(out_path, 'images')
                if os.path.exists(images_path):
                    for img_file in os.listdir(images_path):
                        if img_file.lower().endswith(('.png', '.jpg', '.jpeg')):
                            img_path = os.path.join(images_path, img_file)
                            with open(img_path, 'rb') as img:
                                files = {'file': (img_file, img, 'application/octet-stream')}
                                data = {'bucketname': Config.MINIO_CONFIG['bucket_name']}
                                try:
                                    response = requests.post(
                                        Config.API_CONFIG['upload_url'],
                                        files=files,
                                        data=data
                                    )
                                    if response.status_code != 200:
                                        logger.error(f"上传失败: {img_file}, 状态码: {response.status_code}, 返回: {response.text}")
                                except Exception as e:
                                    logger.error(f"请求错误: {e}")

    def get_all_chunks(self, address: str, dataset_id: str, document_id: str, headers: Dict[str, str], page_size: int = 30) -> List[Dict]:
        """
        获取所有切片信息
        
        Args:
            address: 服务器地址
            dataset_id: 数据集ID
            document_id: 文档ID
            headers: 请求头
            page_size: 每页大小
            
        Returns:
            List[Dict]: 切片列表
        """
        all_chunks = []
        page = 1
        
        while True:
            try:
                # 使用配置中的协议构建URL
                protocol = Config.RAGFLOW_CONFIG.get('protocol', 'http')
                url = f"{protocol}://{address}/api/v1/datasets/{dataset_id}/documents/{document_id}/chunks?page={page}&page_size={page_size}"
                logger.info(f"正在获取第 {page} 页切片...")
                logger.info(f"请求 URL: {url}")
                
                response = requests.get(url, headers=headers)
                logger.info(f"响应状态码: {response.status_code}")
                
                if response.status_code == 200:
                    data = response.json()
                    # chunks = data.get("chunks", [])
                    chunks = data.get("data", {}).get("chunks", [])
                    all_chunks.extend(chunks)
                    logger.info(f"获取到 {len(chunks)} 个切片，累计 {len(all_chunks)} 个切片")
                    
                    # 如果返回的chunks数量小于page_size，说明已经获取完所有数据
                    if len(chunks) < page_size:
                        logger.info("已获取所有切片")
                        break
                    
                    page += 1
                else:
                    error_msg = f"获取切片失败: {response.status_code}"
                    try:
                        # 尝试解析错误响应
                        error_data = response.json()
                        error_msg += f", 错误信息: {error_data.get('message', '未知错误')}, 详情: {error_data}"
                    except json.JSONDecodeError:
                        error_msg += f", 错误内容: {response.text}"
                    
                    logger.error(error_msg)
                    break
            except requests.RequestException as e:
                error_msg = f"网络请求错误: {str(e)}"
                logger.error(error_msg)
                break
            except json.JSONDecodeError as e:
                error_msg = f"JSON解析错误: {str(e)}"
                logger.error(error_msg)
                break
            except Exception as e:
                error_msg = f"未知错误: {str(e)}"
                logger.error(error_msg)
                break
        
        return all_chunks
            

    def get_image_data(self, pdf_name: str) -> Optional[Dict]:
        """
        获取图片数据
        
        Args:
            pdf_name: PDF 文件名
            
        Returns:
            Optional[Dict]: 图片数据，如果文件不存在则返回None
        """
        json_path = os.path.join(
            Config.OUTPUT_DIR,
            pdf_name,
            pdf_name,
            'auto',
            f'{pdf_name}_content_list.json'
        )
        
        if not os.path.exists(json_path):
            logger.error(f"图片数据文件未找到: {json_path}")
            return None

        with open(json_path, 'r', encoding='utf-8') as file:
            return json.load(file)

    # def process_image_data(self, chunks: List[Dict], image_data: Dict, bucket_name: str, image_address_url: str) -> List[Dict]:
    #     """
    #     处理图片数据
        
    #     Args:
    #         chunks: 切片列表
    #         image_data: 图片数据
    #         bucket_name: 存储桶名称
    #         image_address_url: 图片地址URL
            
    #     Returns:
    #         List[Dict]: 处理后的切片列表
    #     """
    #     processed_chunks = []
    #     for chunk in chunks:
    #         original_content = chunk.get('content', '')
    #         new_content = original_content
    
    #         for item in image_data:
    #             if item.get('type') == 'image':
    #                 img_caption = item.get('img_caption', [])
    #                 img_path = item.get('img_path', '')
    
    #                 for caption in img_caption:
    #                     parts = caption.split()
    #                     if len(parts) > 2 and parts[0] == '图':
    #                         figure_name = ' '.join(parts[2:])
    
    #                         if figure_name in original_content:
    #                             img_url = f"{image_address_url}?bucket_name={bucket_name}&file_name={img_path.split('images/', 1)[-1]}"
    #                             image_tag = f'<img src="{img_url}" alt="{figure_name}" style="max-width:100%">'
    #                             new_content = new_content.replace(figure_name, f'{figure_name} {image_tag}')
    
    #         processed_chunks.append({
                
    #             **chunk,
    #             'content': new_content
    #         })
    #     return processed_chunks
  

    def update_chunk(self, address: str, dataset_id: str, document_id: str, chunk_id: str, new_content: str, headers: Dict[str, str]) -> Tuple[bool, Optional[str]]:
        """
        更新切片内容
        
        Args:
            address: 服务器地址
            dataset_id: 数据集ID
            document_id: 文档ID
            chunk_id: 切片ID
            new_content: 新内容
            headers: 请求头
            
        Returns:
            Tuple[bool, Optional[str]]: (是否成功, 失败时的切片ID)
        """
        url = f"http://{address}/api/v1/datasets/{dataset_id}/documents/{document_id}/chunks/{chunk_id}"
        # logger.info(f"正在更新切片 {chunk_id}...")
        # logger.info(f"请求 URL: {url}")
        
        try:
            response = requests.put(
                url,
                headers=headers,
                json={'content': new_content}
            )
            # logger.info(f"响应状态码: {response.status_code}")
            
            if response.status_code == 200:
                # logger.info(f"成功更新切片 {chunk_id}")
                return True, None
            else:
                error_msg = f"更新切片失败: {response.status_code}"
                try:
                    # 尝试解析错误响应
                    error_data = response.json()
                    error_msg += f", 错误信息: {error_data.get('message', '未知错误')}, 详情: {error_data}"
                except json.JSONDecodeError:
                    error_msg += f", 错误内容: {response.text}"
                
                logger.error(error_msg)
                return False, chunk_id
        except requests.RequestException as e:
            error_msg = f"网络请求错误: {str(e)}"
            logger.error(error_msg)
            return False, chunk_id
        except Exception as e:
            error_msg = f"未知错误: {str(e)}"
            logger.error(error_msg)
            return False, chunk_id

    def process_pdf_images(self, input_dir: str, output_dir: str, dataset_id: str, document_id: str) -> Tuple[int, List[str]]:
        """
        处理PDF图片的完整流程
        
        Args:
            input_dir: 输入目录
            output_dir: 输出目录
            dataset_id: 数据集ID
            document_id: 文档ID
            
        Returns:
            Tuple[int, List[str]]: (成功更新的切片数, 失败的切片ID列表)
        """
        # 1. 提取并上传图片
        self.extract_images_from_pdfs(input_dir, output_dir)
        
        # 2. 获取所有数据块
        chunks = self.get_all_chunks

@app.route('/process-pdf', methods=['POST'])
def process_pdf():
    """
    处理PDF文件，提取图片并上传到MinIO
    """
    data = request.json
    input_dir = data.get('input_dir')
    output_dir = data.get('output_dir')
    
    if not input_dir or not output_dir:
        return jsonify({'error': 'input_dir and output_dir are required'}), 400
        
    processor = PDFImageProcessor(api_key='your_api_key')
    processor.extract_images_from_pdfs(input_dir, output_dir)
    
    results = []
    
    # 遍历输出目录中的图片文件
    for root, _, files in os.walk(output_dir):
        for filename in files:
            if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
                img_path = os.path.join(root, filename)
                with open(img_path, 'rb') as img:
                    files = {'file': (filename, img, 'application/octet-stream')}
                    data = {'bucketname': Config.MINIO_CONFIG['bucket_name']}
                    try:
                        response = requests.post(
                            Config.API_CONFIG['upload_url'],
                            files=files,
                            data=data
                        )
                        if response.status_code != 200:
                            logger.error(f"上传失败: {filename}, 状态码: {response.status_code}, 返回: {response.text}")
                        else:
                            results.append({
                                'filename': filename,
                                'status': 'success',
                                'message': response.text
                            })
                    except Exception as e:
                        logger.error(f"请求错误: {e}")
                        results.append({
                            'filename': filename,
                            'status': 'error',
                            'message': str(e)
                        })

    return jsonify({
        'status': 'success',
        'results': results
    })

@app.route('/upload-image', methods=['POST'])
def upload_image():
    """
    上传图片到MinIO
    """
    data = request.json
    output_dir = data.get('output_dir')
    upload_url = data.get('upload_url')
    bucket_name = data.get('bucket_name')
    
    if not output_dir or not upload_url or not bucket_name:
        return jsonify({'error': 'output_dir, upload_url and bucket_name are required'}), 400
        
    # 遍历输出目录中的图片文件
    for root, _, files in os.walk(output_dir):
        for filename in files:
            if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
                img_path = os.path.join(root, filename)
                with open(img_path, 'rb') as img:
                    files = {'file': (filename, img, 'application/octet-stream')}
                    data = {'bucketname': bucket_name}
                    try:
                        response = requests.post(
                            upload_url,
                            files=files,
                            data=data
                        )
                        if response.status_code != 200:
                            logger.error(f"上传失败: {filename}, 状态码: {response.status_code}, 返回: {response.text}")
                    except Exception as e:
                        logger.error(f"请求错误: {e}")

    return jsonify({
        'status': 'success'
    })

if __name__ == '__main__':
    app.run(
        host='0.0.0.0',
        port=18005,
        debug=True
    )
